import asyncio
import logging
import os
import sys
import time

import dns.asyncresolver
from dns import message
from dns.rdataclass import RdataClass
from dns.rdatatype import RdataType
from motor import MotorClient
from motor.core import AgnosticCollection

from transform import timestampToTime

logging.basicConfig(format='[%(asctime)s]{%(filename)s:%(lineno)d}[%(levelname)s] %(message)s', level=logging.DEBUG)
logger = logging.getLogger(__name__)

base_nameservers = ['8.8.8.8']
dns.resolver.default_resolver = dns.asyncresolver.Resolver(configure=False)
dns.resolver.default_resolver.nameservers = base_nameservers
client = MotorClient('mongodb://lxt:aa347347@173.82.107.218:27017/?authSource=admin')
db = client['dns']
collection: AgnosticCollection = db['resolvers']
collection_timeout: AgnosticCollection = db['resolvers_info']

collection_domains: AgnosticCollection = db['domains']

total_ip = 0
MAX_PORTS = 10000
sem = asyncio.Semaphore(MAX_PORTS)
MAX_INTERVAL = 1000


async def getTTL(domain: str):
    resolver = dns.asyncresolver.Resolver(configure=False)
    resolver.nameservers = base_nameservers

    name = dns.name.from_text(domain)
    logger.info(f'query name: {name}')
    res = await resolver.resolve(name, 'A')

    # Parse cname chain
    while True:
        cnameSet = res.response.get_rrset(message.ANSWER, name, rdclass=RdataClass.IN, rdtype=RdataType.CNAME)
        if not cnameSet:
            break
        cname: dns.rdtypes.ANY.CNAME.CNAME = cnameSet.pop()
        name = dns.name.from_text(cname.to_text())

    parent = name.parent()  # find the NS of its parent zone
    # Find authority NS, from bottom to top, iteratively
    while True:
        try:
            res = await resolver.resolve(parent, rdtype=RdataType.NS)
            break
        except dns.resolver.NXDOMAIN:
            parent = parent.parent()
        except dns.resolver.NoAnswer:
            parent = parent.parent()

    nsSet = res.response.get_rrset(message.ANSWER, parent, RdataClass.IN, RdataType.NS)
    ns: dns.rdtypes.ANY.NS.NS = nsSet.pop()  # randomly select a NS server

    # query the IP of the NS
    nsName = dns.name.from_text(ns.to_text())
    res = await resolver.resolve(nsName)
    nsSet = res.response.get_rrset(message.ANSWER, nsName, RdataClass.IN, RdataType.A)
    # print(nsSet)

    # query the authority NS for the TTL of the A record
    resolver.nameservers = [ns.to_text() for ns in nsSet]
    res = await resolver.resolve(name, raise_on_no_answer=False)
    if not res.response.answer:
        assert res.response.authority, 'No NS record found'
        authSet = res.response.authority.pop()
        nsIPs = []
        if authSet:
            for ns in authSet:
                nsIp = res.response.get_rrset(message.ADDITIONAL, dns.name.from_text(ns.to_text()), RdataClass.IN,
                                              RdataType.A)
                if nsIp:
                    nsIPs.append(nsIp.pop())

        if not nsIPs:
            nsName = ns.to_text()
            print(f'Getting NS IP: {nsName}')
            res = await dns.resolver.resolve(nsName)
            nsIp = res.response.get_rrset(message.ANSWER, dns.name.from_text(ns.to_text()), RdataClass.IN,
                                          RdataType.A).pop()
            print(f'nsIP: {nsIp}')
            nsIPs.append(nsIp)
            time.sleep(1)

        resolver.nameservers = [nsIPs[-1].to_text()]
        res = await resolver.resolve(name)

    print(f'authority NS: {ns}')

    rrSet = res.response.get_rrset(message.ANSWER, name, RdataClass.IN, RdataType.A)
    if rrSet is not None:
        ip: dns.rdtypes.IN.A.A
        As = [ip.to_text() for ip in rrSet]
    else:
        As = []

    print(f'Real name: {name}')
    print(f'Real A: {As}')
    print(f'ttl: {rrSet.ttl}')
    return {
        "name": domain,
        "realName": name.to_text(),
        "A": As,
        "ttl": rrSet.ttl
    }


async def saveAuthorityTTLs(filename: str, limit=-1):
    with open(filename, 'r', encoding='utf-8') as f:
        lines = f.readlines()[:limit] if limit > 0 else f.readlines()
        coro = [saveOneDomain(line) for line in lines]
        await asyncio.gather(*coro, return_exceptions=True)


async def saveOneDomain(line):
    try:
        name, domain = line.rstrip().split(' ')
        res = await getTTL(domain)
        res['cnName'] = name
        await collection_domains.update_one({'name': res['name']}, {'$set': res}, upsert=True)
    except Exception as e:
        print(e)


async def getRRSet(domain: str, nameserver: str):
    domain = dns.name.from_text(domain)
    resolver = dns.asyncresolver.Resolver(configure=False)
    resolver.nameservers = [nameserver]

    query_time = time.time()
    try:
        res = await resolver.resolve(domain)
    except dns.resolver.LifetimeTimeout as e:
        # logger.debug(f'time out: {nameserver}')
        async with await client.start_session() as s:
            await collection_timeout.update_one({
                'ip': nameserver
            }, {
                '$set': {
                    'available': False,
                    'reason': str(e),
                    'updateTime': timestampToTime(time.time())
                }
            }, upsert=True, session=s)
        return
    rrSet = res.response.get_rrset(message.ANSWER, domain, RdataClass.IN, RdataType.A)
    return query_time, rrSet


async def probeQueryRate(domain: str, name: str, totalTTL: int, nameserver: str, comment: str, times: int = 5):
    global total_ip
    total_ip += 1

    window = totalTTL / 10
    pre_tr = None
    try:
        for i in range(times):
            async with sem:
                print(domain, nameserver)
                query_time_rrSet = await getRRSet(domain, nameserver)

            if query_time_rrSet is None:
                return
            query_time, rrSet = query_time_rrSet
            remain_ttl = rrSet.ttl
            tr = query_time - (totalTTL - remain_ttl)

            logger.debug(f'[{comment}][#{i + 1}]query time(Tp): {query_time}')
            logger.debug(f'[{comment}][#{i + 1}]remain ttl(Tl): {remain_ttl}')
            logger.debug(f'[{comment}][#{i + 1}]tr: {tr}')

            if remain_ttl == totalTTL:
                logger.info(f'[{comment}][#{i + 1}]Cache was refreshed by us')
                delta_tr = window
            else:
                delta_tr = (tr - pre_tr - totalTTL) if pre_tr else None

            interval = window + remain_ttl
            if interval > window + totalTTL:
                logger.error(f'interval: {interval} is larger than '
                             f'window + totalTTL: {window + totalTTL}, skip probing')
                return

            # 指数衰减探测
            if delta_tr is not None and delta_tr < 0:
                # if interval > 2:  # minimum interval
                #     interval /= 2
                # i -= 1
                # logger.warning(f'[{comment}][#{i+1}]delta_tr < 0, start log probe, ip: {nameserver}, '
                #                f'current interval: {interval}')
                logger.warning(f'[{comment}][#{i + 1}]negative delta_tr, sleep for interval: {interval}')
                await asyncio.sleep(interval)

            else:
                # restore interval
                # interval = totalTTL
                logger.info(f'[{comment}][#{i + 1}]delta_tr: {delta_tr}')
                if delta_tr is not None:
                    ip: dns.rdtypes.IN.A.A
                    As = [ip.to_text() for ip in rrSet]
                    async with await client.start_session() as s:
                        await collection.insert_one({
                            'ip': nameserver,
                            'domain': domain,
                            'cn_name': name,
                            'A': As,
                            'ttl': totalTTL,
                            'tr': tr,
                            'pre_tr': pre_tr,
                            'delta_tr': delta_tr,
                            'query_time': query_time,
                            'remain_ttl': remain_ttl,
                            'interval': interval,
                            'window': window,
                            'updateTime': timestampToTime(time.time())
                        }, session=s)

                await asyncio.sleep(interval)
            pre_tr = tr

    except Exception as e:
        logger.warning(f'nameserver: {nameserver} failed: {e}')
        try:
            async with await client.start_session() as s:
                await collection_timeout.update_one({
                    'ip': nameserver
                }, {
                    '$set': {
                        'available': False,
                        'reason': str(e),
                        'updateTime': timestampToTime(time.time())
                    }
                }, upsert=True, session=s)
        except:
            pass
        return


async def main():
    records = collection_domains.find()
    with open('data/resolverIps.txt', 'r') as f:
        ips = list(map(str.rstrip, f.readlines()))

    batch = 1
    tasks = []
    count = 1
    async for record in records:
        realName = record['realName']
        totalTTL = record['ttl']
        if totalTTL > MAX_INTERVAL:
            continue

        name = record['cnName']
        for i in range(len(ips)):
            if count % 10000 == 0:
                count = 1
                batch += 1
                await asyncio.gather(*tasks)
                tasks.clear()

            logger.debug(f'creating coroutine: #{count}')
            task = asyncio.create_task(probeQueryRate(realName, name, totalTTL, ips[i], f'{batch}:{name}', 10))
            tasks.append(task)
            count += 1

    await asyncio.gather(*tasks)

if __name__ == '__main__':
    # domains = ["news.qq.com", "www.baidu.com", "www.taobao.com", "www.sina.com.cn", "weibo.com"]

    # print(getTr(domains[0]))
    if sys.platform != "win32":
        import uvloop

        uvloop.install()
        os.system(f'ulimit -n {MAX_PORTS + 100}')

    try:
        # asyncio.run(main())
        asyncio.run(saveAuthorityTTLs('data/topDomains.txt'))
    finally:
        print('over')

    # try:
    #     asyncio.run(saveAuthorityTTLs('data/topDomains.txt'))

    # finally:
    #     print('over')
